#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Created by Ross on 18-10-14
import os

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.utils import to_categorical

import utils
from model.LMTGRU import CNNLMTGRU
from utils import generate_batch

DATA_PATH = 'data/smp2018/10flod_all_3600+_未转阿拉伯'
EPOCHS = 1000
LOG_DIR = 'LMTGRU_train_log'
TRAIN_DIR = 'LMTGRU_saved_model'

test_x = np.load(os.path.join(DATA_PATH, 'test_x.npy'))
test_y = np.load(os.path.join(DATA_PATH, 'test_y.npy')).astype(np.int32)
test_x = pad_sequences(test_x, 20, 'float32')
# test_y = np.where(test_y == 3, 1, 0)
# test_y = to_categorical(test_y, 2)
test_y = to_categorical(test_y, 31)

for fold in range(10):
    train_x, train_y, dev_x, dev_y = utils.get_single_fold(DATA_PATH, fold)
    train_x = pad_sequences(train_x, 20, dtype='float32')
    dev_x = pad_sequences(dev_x, 20, dtype='float32')

    # train_y = np.where(train_y == 3, 1, 0)
    # dev_y = np.where(dev_y == 3, 1, 0)
    train_y = to_categorical(train_y, 31)
    dev_y = to_categorical(dev_y, 31)
    # print(train_y.shape)
    # train_y = to_categorical(train_y, 2)
    # dev_y = to_categorical(dev_y, 2)
    tf.reset_default_graph()
    model = CNNLMTGRU(256, [3, 4, 5], 2, 250, 20, 400, 0.001)
    with tf.Session() as sess:
        # sess = tf_debug.LocalCLIDebugWrapperSession(sess)
        summary = tf.summary.FileWriter(os.path.join(LOG_DIR, str(fold)), sess.graph)
        max_dev_acc = model.start_or_continue_training(sess, os.path.join(TRAIN_DIR, str(fold)))
        for epoch in range(1, EPOCHS + 1):
            for x, y in generate_batch(train_x, train_y, 64, shuffle=True, undersampling=False):
                model.train(sess, x, y, summary)
            if epoch % 2 == 0:
                loss = model.compute_loss(sess, train_x, train_y)
                acc = model.compute_accuracy(sess, train_x, train_y)
                print('train acc, loss:', acc, loss)

                loss = model.compute_loss(sess, dev_x, dev_y)
                dev_acc = model.compute_accuracy(sess, dev_x, dev_y)
                print('dev acc, loss: {}, {}'.format(dev_acc, loss))
                # pred = model.predict(sess, dev_x)
                # df = pd.DataFrame(list(zip(np.squeeze(pred), np.squeeze(dev_y))), columns=['pred', 'y'])
                # df.to_csv('pred.csv', index=False)
                if dev_acc > max_dev_acc:
                    loss = model.compute_loss(sess, test_x, test_y)
                    acc = model.compute_accuracy(sess, test_x, test_y)
                    pred_test = model.predict(sess, test_x)
                    print('test acc, loss: {}, {}'.format(acc, loss))

                    max_dev_acc = dev_acc
                    print('max:', max_dev_acc)
                    model.save(sess, max_dev_acc, os.path.join(TRAIN_DIR, str(fold)))

                    df = pd.DataFrame(list(zip(pred_test, test_y)), columns=['pred', 'ground_true'])
                    df.to_csv('LMTGRUpred.csv', index=False)
                print()

if __name__ == '__main__':
    pass
